//
// Created by 12992 on 24-10-28.
//

#include "include/args.h"
#include "crypto/base64.h"
#include "crypto/rsa.h"
#include <fstream>
using namespace std;

args::args()
    : program("ex-openssl"), base64("base64"), rsa("rsa"),
      encode("encode"), decode("decode"),
      gen_key_pair("generate"), decrypt("decrypt"), encrypt("encrypt")
{
    this->program.add_subparser(this->base64);
    this->program.add_subparser(this->rsa);
    this->init_rsa();
    this->init_base64();
}

void args::parse(int argc, const char** argv)
{
    try
    {
        this->program.parse_args(argc, argv);

        if (this->program.is_subcommand_used(this->base64))
        {
            // 使用base64
            if (this->base64.is_subcommand_used(this->encode))
            {
                // base64编码操作
                this->base64_encode();
            }
            else if (this->base64.is_subcommand_used(this->decode))
            {
                // base64解码操作
                this->base64_decode();
            }
        }
        else if (this->program.is_subcommand_used(this->rsa))
        {
            // 使用rsa
            if (this->rsa.is_subcommand_used(this->encrypt))
            {
                // rsa 加密操作
                this->rsa_encrypt();
            }
            else if (this->rsa.is_subcommand_used(this->decrypt))
            {
                // rsa 解密操作
                this->rsa_decrypt();
            }
            else if (this->rsa.is_subcommand_used(this->gen_key_pair))
            {
                // 生成密钥
                this->generate_key();
            }
        }
    }
    catch (const exception& e)
    {
        cerr << e.what() << endl;
        exit(1);
    }
}

void add_input_output_file(argparse::ArgumentParser& parser)
{
    parser.add_argument("-i", "--input").help("input file").metavar("<INPUT_PATH>").required();
    parser.add_argument("-o", "--output").help("output file").metavar("<OUTPUT_PATH>").required();
}

void args::init_rsa()
{
    this->rsa.add_description("Perform encryption and decryption using RSA.");
    this->rsa.add_subparser(this->gen_key_pair);
    this->rsa.add_subparser(this->encrypt);
    this->rsa.add_subparser(this->decrypt);
    this->gen_key_pair.add_description("Generate rsa key pair.");
    this->decrypt.add_description("Rsa decrypt");
    this->encrypt.add_description("Rsa encrypt");
    this->decrypt.add_argument("--public-key").help("public key path").metavar("PATH");
    this->decrypt.add_argument("--private-key").help("private key path").metavar("PATH");
    this->encrypt.add_argument("--public-key").help("public key path").metavar("PATH");
    this->encrypt.add_argument("--private-key").help("private key path").metavar("PATH");
    add_input_output_file(this->decrypt);
    add_input_output_file(this->encrypt);
}


void args::init_base64()
{
    this->base64.add_description("Encode and decode using base64.");
    this->base64.add_subparser(this->encode);
    this->base64.add_subparser(this->decode);
    this->encode.add_description("Base64 Encode");
    this->decode.add_description("Base64 Decode");
    add_input_output_file(this->encode);
    add_input_output_file(this->decode);
}

void args::base64_encode()
{
    auto input_path = this->encode.get<string>("--input");
    auto output_path = this->encode.get<string>("--output");
    ifstream input(input_path, ios::binary);
    ofstream output(output_path);
    if (!input.is_open())
    {
        throw runtime_error("Input file not found");
    }
    if (!output.is_open())
    {
        throw runtime_error("Can not open Output file");
    }
    input.seekg(0, ios::end);
    auto size = input.tellg();
    input.seekg(0, ios::beg);
    string buffer(size, '\0');
    input.read(&buffer[0], size);
    crypto::base64 base64;
    auto output_buffer = base64.encrypt(buffer);
    output.write(output_buffer.c_str(), output_buffer.size());
}

void args::base64_decode()
{
    auto input_path = this->decode.get<string>("--input");
    auto output_path = this->decode.get<string>("--output");
    ifstream input(input_path, ios::binary);
    ofstream output(output_path);
    if (!input.is_open())
    {
        throw runtime_error("Input file not found");
    }
    if (!output.is_open())
    {
        throw runtime_error("Can not open Output file");
    }
    input.seekg(0, ios::end);
    auto size = input.tellg();
    input.seekg(0, ios::beg);
    string buffer(size, '\0');
    input.read(&buffer[0], size);
    crypto::base64 base64;
    auto output_buffer = base64.decrypt(buffer);
    output.write(output_buffer.c_str(), output_buffer.size());
}

void args::generate_key()
{
    crypto::rsa::generate_key_pair("private.pem", "public.pem");
}

void args::rsa_decrypt()
{
    bool use_private = this->decrypt.is_used("--private-key");
    bool use_public = this->decrypt.is_used("--public-key");
    auto input_path = this->decrypt.get<string>("--input");
    auto output_path = this->decrypt.get<string>("--output");
    string key_path;
    ifstream input(input_path, ios::binary);
    ofstream output(output_path, ios::binary);
    if (!input.is_open())
    {
        throw runtime_error("Input file not found");
    }
    if (!output.is_open())
    {
        throw runtime_error("Can not open Output file");
    }
    input.seekg(0, ios::end);
    auto size = input.tellg();
    input.seekg(0, ios::beg);
    string buffer(size, '\0');
    input.read(&buffer[0], size);
    if (use_private && use_public)
    {
        throw runtime_error("Cannot provide both the private key and the public key simultaneously.");
    }
    else if (use_public)
    {
        key_path = this->decrypt.get<string>("--public-key");
        this->rsa_public_decrypt(output, key_path, buffer);
    }
    else if (use_private)
    {
        key_path = this->decrypt.get<string>("--private-key");
        this->rsa_private_decrypt(output, key_path, buffer);
    }
    else
    {
        throw runtime_error("Need to provide either the private key or the public key.");
    }
}

void args::rsa_encrypt()
{
    bool use_private = this->encrypt.is_used("--private-key");
    bool use_public = this->encrypt.is_used("--public-key");
    auto input_path = this->encrypt.get<string>("--input");
    auto output_path = this->encrypt.get<string>("--output");
    string key_path;
    ifstream input(input_path, ios::binary);
    ofstream output(output_path, ios::binary);
    if (!input.is_open())
    {
        throw runtime_error("Input file not found");
    }
    if (!output.is_open())
    {
        throw runtime_error("Can not open Output file");
    }
    input.seekg(0, ios::end);
    auto size = input.tellg();
    input.seekg(0, ios::beg);
    string buffer(size, '\0');
    input.read(&buffer[0], size);
    if (use_private && use_public)
    {
        throw runtime_error("Cannot provide both the private key and the public key simultaneously.");
    }
    else if (use_public)
    {
        key_path = this->encrypt.get<string>("--public-key");
        this->rsa_public_encrypt(output, key_path, buffer);
    }
    else if (use_private)
    {
        key_path = this->encrypt.get<string>("--private-key");
        this->rsa_private_encrypt(output, key_path, buffer);
    }
    else
    {
        throw runtime_error("Need to provide either the private key or the public key.");
    }
}

void args::rsa_public_encrypt(ofstream& output, const string& key_path, const string& data)
{
    crypto::rsa rsa(crypto::rsa::key_type::PUBLIC, key_path);
    auto res = rsa.encrypt(data);
    output.write(res.c_str(), res.size());
}

void args::rsa_private_encrypt(ofstream& output, const string& key_path, const string& data)
{
    crypto::rsa rsa(crypto::rsa::key_type::PRIVATE, key_path);
    auto res = rsa.encrypt(data);
    output.write(res.c_str(), res.size());
}

void args::rsa_public_decrypt(std::ofstream& output, const std::string& key_path, const std::string& data)
{
    crypto::rsa rsa(crypto::rsa::key_type::PUBLIC, key_path);
    auto res = rsa.decrypt(data);
    output.write(res.c_str(), res.size());
}

void args::rsa_private_decrypt(std::ofstream& output, const std::string& key_path, const std::string& data)
{
    crypto::rsa rsa(crypto::rsa::key_type::PRIVATE, key_path);
    auto res = rsa.decrypt(data);
    output.write(res.c_str(), res.size());
}
